import os
import subprocess
import sys
import tensorflow as tf
import shutil
import datetime


exe_steps = []
exe_steps.append('prepare datas')
exe_steps.append('start RGB opt')
exe_steps.append('step0: start unwrap')
exe_steps.append('step1: start fit AlbedoNormal_RPB')
exe_steps.append('step2: generate tex')
exe_steps.append('step2: generate norm')
exe_steps.append('step3: convert_texture_domain')
exe_steps.append('output results')

class timecounter():
    def __init__(self,name):        
        self.start_time = datetime.datetime.now()
        self.name = name
    def timediff(self):
        end_time = datetime.datetime.now()
        time_diff = (end_time - self.start_time)
        execution_time = time_diff.total_seconds() * 1000
        print("{} use time {} ms".format(self.name,execution_time))
    

def print_time(content):
    current_time = datetime.datetime.now()
    formatted_time = current_time.strftime("%Y-%m-%d %H:%M:%S")

    # 打印时间字符串
    print("[{}]: {}".format(formatted_time,content))

GPU_NO=0
is_bfm=False

pwd = os.path.dirname(os.path.abspath(__file__))
pwd = pwd.replace('\\','/')
sys.path.append(pwd)



# # constants
basic_path=pwd + '/3DMM/Tencent_AI-NEXT_3D_Face_Model_East-Asian_Version/files'
resources_path=pwd + '/3DMM/resources/resources'

uv_base=basic_path +"/AI-NEXT-Albedo-Global.mat"
uv_regional_pyramid_base=basic_path + "/AI-NEXT-AlbedoNormal-RPB"

if is_bfm == False:
    shape_exp_bases=basic_path + "/AI-NEXT-Shape-NoAug.mat"
else:
    shape_exp_bases=resources_path + "/BFM2009_Model.mat"


vggpath=resources_path + "/vgg-face.mat"
pb_path=resources_path + "/PB"

# # data directories
# ROOT_DIR=pwd + "/3DMM/test_data/test_data/RGB/test1/single_img"
ROOT_DIR=pwd + "/3DMM/test_data/test_data/RGB/test1/single_gr"
img_dir=ROOT_DIR


totalTime = timecounter('totalTime')
########################################################
print_time("prepare datas")


prepare_dir=ROOT_DIR + "/prepare_rgb"


command = "python -u ./data_prepare/run_data_preparation.py --GPU_NO={}  --mode=test_RGB --pb_path={} --img_dir={} --out_dir={}".format(GPU_NO,pb_path,img_dir,prepare_dir)

if "prepare datas" in exe_steps:
    
    t = timecounter('prepare datas')
    result = subprocess.check_output(command).decode()
    print_time(result)
    t.timediff()
print_time("prepare datas done")



########################################################
print_time("start RGB opt")

if is_bfm == False:
    shape_out_dir=ROOT_DIR + "/our_opt_RGB"
else:
    shape_out_dir=ROOT_DIR + "/bfm_opt_RGB"



train_step=150
log_step=20
learning_rate=0.05
lr_decay_step=20
lr_decay_rate=0.9

photo_weight=100.0
gray_photo_weight=80.0
reg_shape_weight=0.5
reg_tex_weight=2.0
id_weight=1.0
real_86pt_lmk3d_weight=5.0
real_68pt_lmk2d_weight=5.0
lmk_struct_weight=0

num_of_img=1
project_type="Pers"

command = "python ./optimization/rgb/run_RGB_opt.py " +\
"--GPU_NO={} ".format(GPU_NO) +\
"--is_bfm={} ".format(is_bfm) +\
"--basis3dmm_path={} ".format(shape_exp_bases) +\
"--uv_path={} ".format(uv_base) +\
"--vggpath={} ".format(vggpath) +\
"--base_dir={} ".format(prepare_dir) +\
"--log_step={} ".format(log_step) +\
"--train_step={} ".format(train_step) +\
"--learning_rate={} ".format(learning_rate) +\
"--lr_decay_step={} ".format(lr_decay_step) +\
"--lr_decay_rate={} ".format(lr_decay_rate) +\
"--photo_weight={} ".format(photo_weight) +\
"--gray_photo_weight={} ".format(gray_photo_weight) +\
"--id_weight={} ".format(id_weight) +\
"--reg_shape_weight={} ".format(reg_shape_weight) +\
"--reg_tex_weight={} ".format(reg_tex_weight) +\
"--real_86pt_lmk3d_weight={} ".format(real_86pt_lmk3d_weight) +\
"--real_68pt_lmk2d_weight={} ".format(real_68pt_lmk2d_weight) +\
"--lmk_struct_weight={} ".format(lmk_struct_weight) +\
"--num_of_img={} ".format(num_of_img) +\
"--out_dir={} ".format(shape_out_dir) +\
"--summary_dir={}/summary ".format(shape_out_dir) +\
"--project_type={} ".format(project_type)


if "start RGB opt" in exe_steps:
    
    t = timecounter('start RGB opt')
    result = subprocess.check_output(command).decode()
    print_time(result)
    t.timediff()
print_time("start RGB opt done")


########################################################
if is_bfm == False:
    print_time("start generate HD texture")    

    print_time("step0: start unwrap")
    os.environ['CUDA_VISIBLE_DEVICES']=str(GPU_NO)
    
    command = "python -u ./texture/step0_unwrapper.py " +\
        "--basis3dmm_path={} ".format(shape_exp_bases) +\
        "--resources_path={} ".format(resources_path) +\
        "--uv_path={} ".format(uv_base) +\
        "--uv_size=512 " +\
        "--is_mult_view=False " +\
        "--is_orig_img=True " +\
        "--input_dir={} ".format(shape_out_dir) +\
        "--output_dir={}/unwrap ".format(ROOT_DIR)

    if "step0: start unwrap" in exe_steps:
        t = timecounter('step0: start unwrap')
        result = subprocess.check_output(command).decode()
        print_time(result)
        t.timediff()
    print_time("step0: start unwrap done")

    print_time("step1: start fit AlbedoNormal_RPB")
    
    uv_weight_mask_path = resources_path + '/uv_boundary_masks/boundary_512.png'
    os.environ['CUDA_VISIBLE_DEVICES']=str(GPU_NO)
    command = "python -u ./texture/step1_fit_AlbedoNormal_RPB.py " +\
        "--basis3dmm_path={} ".format(shape_exp_bases) +\
        "--resources_path={} ".format(resources_path) +\
        "--uv_path={} ".format(uv_regional_pyramid_base) +\
        "--write_graph=False " +\
        "--data_dir={}/unwrap ".format(ROOT_DIR) +\
        "--out_dir={}/fit_unwrap ".format(ROOT_DIR) +\
        "--uv_tv_weight=0.1 " +\
        "--uv_reg_tex_weight=0.001 " +\
        "--learning_rate=0.1 " +\
        "--train_step=200 " +\
        "--uv_weight_mask_path={} ".format(uv_weight_mask_path)

    if "step1: start fit AlbedoNormal_RPB" in exe_steps:
        
        t = timecounter('step1: start fit AlbedoNormal_RPB')
        result = subprocess.check_output(command).decode()
        print_time(result)
        t.timediff()
    print_time("step1: start fit AlbedoNormal_RPB done")

    print_time("step2: generate tex")
    
    os.environ['CUDA_VISIBLE_DEVICES']=str(GPU_NO)
    command = "python -u ./texture/step2_pix2pix.py --mode texture --func test --pb_path {}/pix2pix_tex.pb \
        --input_dir={}/fit_unwrap \
        --output_dir={}/pix2pix \
        --resources_path={} ".format(pb_path,ROOT_DIR,ROOT_DIR,resources_path)

   
    if "step2: generate tex" in exe_steps:
        t = timecounter('step2: generate tex')
        result = subprocess.check_output(command).decode()
        print_time(result)
        t.timediff()
    print_time("step2: generate tex done")



    print_time("step2: generate norm")
    
    os.environ['CUDA_VISIBLE_DEVICES']=str(GPU_NO)
    command = "python -u ./texture/step2_pix2pix.py --mode normal --func test --pb_path {}/pix2pix_norm.pb \
        --input_dir={}/fit_unwrap \
        --output_dir={}/pix2pix \
        --resources_path={} ".format(pb_path,ROOT_DIR,ROOT_DIR,resources_path)

    if "step2: generate norm" in exe_steps:
        t = timecounter('step2: generate norm')
        result = subprocess.check_output(command).decode()
        print_time(result)
        t.timediff()
    print_time("step2: generate norm done")


    FIT_DIR=ROOT_DIR + "/fit_unwrap"
    PIX_DIR=ROOT_DIR + "/pix2pix"
    UNWRAP_DIR=ROOT_DIR + "/unwrap"
    OUT_DIR=ROOT_DIR + "/pix2pix_convert"

    print_time("step3: convert_texture_domain")
    command = "python -u ./texture/step3_convert_texture_domain.py " +\
        "--input_fit_dir={} ".format(FIT_DIR) +\
        "--input_pix2pix_dir={} ".format(PIX_DIR) +\
        "--input_unwrap_dir={} ".format(UNWRAP_DIR) +\
        "--output_dir={} ".format(OUT_DIR) +\
        "--resources_path={} ".format(resources_path) 

    if "step3: convert_texture_domain" in exe_steps:
        t = timecounter('step3: convert_texture_domain')
        result = subprocess.check_output(command).decode()
        print_time(result)
        t.timediff()
    print_time("step3: convert_texture_domain done")




# ########################################################
print_time("output results")

results_dir=ROOT_DIR + "/results"
if not os.path.exists(results_dir):
    os.makedirs(results_dir)
    

if is_bfm == False:
    shutil.copy(pwd + "/test.mtl", results_dir)
    shutil.copy(shape_out_dir + "/face.obj", results_dir + "/head.obj")
    shutil.copy(ROOT_DIR + "/pix2pix_convert/output_for_texture_tex_D.png", results_dir + "/albedo.png")
    shutil.copy(ROOT_DIR + "/pix2pix/out_for_texture_tex_N.png", results_dir + "/normal.png")
else:
    shutil.copy(shape_out_dir + "/face.ply","results_dir/")


print_time("output results done")
print_time("finish")
totalTime.timediff()